feat: add SQLCommenter observability to database drivers (fixes #2899) by tanujbolisetty · Pull Request #2913 · googleapis/mcp-toolbox

tanujbolisetty · 2026-03-31T19:28:47Z

🚀 Overview

This PR introduces SQLCommenter observability across all supported database drivers in the GenAI Toolbox. It automatically injects rich contextual metadata (tool name, framework, agent controller, database driver, and OpenTelemetry traceparent) into every query executed by the toolbox.

By standardizing this telemetry pattern, downstream database administrators and cloud monitoring tools (like Cloud SQL Insights) can seamlessly map database loads back to specific AI agents and MCP tools without any manual parsing.

🛠 Features & Modifications

1. Comprehensive Driver Coverage

We have instrumented the RunSQL method on 100% of SQL-capable database sources.
The metadata is injected safely as a block comment (/* ... */) before the trailing semicolon (where applicable) to preserve syntax parsing.

Instrumented Standard SQL Drivers: alloydbpg, clickhouse, cockroachdb, cloudsqlmssql, cloudsqlmysql, cloudsqlpg, firebird, mindsdb, mssql, mysql, oceanbase, oracle, postgres, singlestore, snowflake, spanner, sqlite, tidb, trino, yugabytedb.
Instrumented Query/Graph Drivers: cassandra (CQL) and couchbase (N1QL). For Couchbase, the Source RunSQL signature and corresponding Tool wrapper Invoke interface were refactored to explicitly accept context.Context to enable propagation.
Native Labeling: bigquery is instrumented via native BigQuery Job Labels instead of SQL Comments, leveraging the same contextual map.
Exclusions (Safety): Object/NoSQL stores (mongodb, firestore, redis, bigtable, dgraph) were explicitly skipped as they do not support standard SQL comment blocks, preventing catastrophic parsing crashes.

2. Context Injection Pipeline

Created internal/sqlcommenter/commenter.go to handle context reading, deterministic key sorting, formatting (RFC 3986 URL encoding), and OpenTelemetry span extraction.
MCP Auto-Context Pipeline: internal/server/server.go was updated to trap the clientInfo.name payload from the MCP initialize handshake in a concurrent mcpClientNames sync map. api.go extracts this dynamically during tool invocation to populate the controller field automatically (e.g. associating queries directly with "Mosa" or "Claude Desktop").
Fallback configurations TOOLBOX_APP_NAME and TOOLBOX_AGENT_NAME were established to manually override the context if needed.

3. Core Framework & Build Architecture Upgrades

These changes resolve significant localized testing and binary compilation blockers encountered during development:

API UI Autostart: Modified internal/server/server.go configuration booting logic so that passing the --ui config implicitly turns on the REST API endpoints (if cfg.EnableAPI || cfg.UI).
Oracle CGO Decoupling: Separated the Oracle driver implementations into two new files (oci_cgo.go and oci_pure.go). By leveraging Go build tags (//go:build cgo), the toolbox now dynamically compiles using the lightweight go-ora/v2 Pure Go driver by default, entirely sidestepping notoriously complex godror C compiler dependencies native to Oracle cross-compilation.

🧪 Verification

Full test matrix suite completely passing: go test ./internal/sqlcommenter/....
Tested the append function string manipulation algorithm safely by mocking the WithDBDriver and WithToolName context decorators directly, avoiding the need for live mocked database instances in CI.
go build ./... successfully compiles across pure GO targets natively.

📝 Documentation

Overhauled docs/SQLCOMMENTER_README.md and explicitly mapped out the full driver support matrix, environment overrides, and NoSQL exclusions.
Recommendation: Ensure DEVELOPER.md or the primary README.md is updated to highlight that the default codebase cross-compiles with Oracle's PureGo driver to unblock new developers!

google-cla · 2026-03-31T19:28:57Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

gemini-code-assist

Code Review

This pull request implements the SQLCommenter observability layer across all SQL-capable database drivers, enabling query tracking by tool name, agent identity, and trace metadata. It introduces the internal/sqlcommenter package and updates the MCP server to persist agent names across sessions. Feedback identifies a memory leak in the session name map, suggests refactoring duplicated RunSQL logic in the CockroachDB source, and recommends using %20 for space encoding to ensure strict adherence to the SQLCommenter specification.

…eapis#2899)

…Manager

…bstraction drivers

drstrangelooker · 2026-04-10T15:20:30Z

Can you please integrate the latest changes on main so we can see what is actually changing?

…ramework defaults

tanujbolisetty · 2026-04-10T16:28:45Z

"I've just successfully rebased onto the latest upstream/main (v1.0.0 mcp-toolbox namespace) and squashed the history. The PR is now completely up-to-date with a clean commit trail!"

Yuan325

Hi @tanujbolisetty Thank you for your contributions! Please let me know if you need any clarifications. :)

Yuan325 · 2026-04-14T18:02:12Z

These changes should be reversed. We did not update the llamaindex sdk repo name~

Yuan325 · 2026-04-14T18:03:02Z

+- "Execute a DML query to update customer names."
+- "List all tables in the `my-database`."
+- "Execute a DQL query to select data from `orders` table."


Let's keep to * to be consistent with other readmes.

Yuan325 · 2026-04-14T18:03:37Z

Changes in this file should be reversed. It is redundant for this sqlcommenter.

Yuan325 · 2026-04-14T18:08:09Z

This doc should not be included here. We can add sqlcommenter under docs/en/documentation/monitoring. Please modify this file to fit with the docs template format. Alternatively, since every source's sqlcommenter is different, we can also add it under each docs/en/integrations source's docs.

Yuan325 · 2026-04-14T18:10:08Z

+| **CloudSQL Postgres** | `internal/sources/cloudsqlpg/cloud_sql_pg.go` | SQL Comment (`/* ... */`) |
+| **CockroachDB** | `internal/sources/cockroachdb/cockroachdb.go` | SQL Comment (`/* ... */`) |
+| **Couchbase (N1QL)** | `internal/sources/couchbase/couchbase.go`<br>`internal/tools/couchbase/couchbase.go` | N1QL Comment (`/* ... */`) |
+


Suggested change

Yuan325 · 2026-04-14T18:33:15Z

let's remove this file. All docs should go under docs. We can concatenate the information here along with the other sqlcommenter documentation.

Yuan325 · 2026-04-14T18:34:33Z

+	// Embed the tool name in the context so RunSQL can attach it as a BigQuery
+	// Job Label for Cloud SQL Insights / BigQuery observability.
+	ctx = sqlcommenter.WithToolName(ctx, t.Name)
+


This should be added in server/mcp?

Yuan325 · 2026-04-14T18:35:04Z

Please reverse, this change is redundant~

Yuan325 · 2026-04-14T18:35:31Z

+func (s *Source) RunSQL(ctx context.Context, statement string, params parameters.ParamValues, isQuery bool, timeout string) (any, error) {
+	// Dgraph does not support SQL block comments, so we explicitly skip injecting SQLCommenter metadata here.
+


since we are not using ctx here, please reverse this.

Yuan325 · 2026-04-14T18:35:44Z

redundant, please reverse~

tanujbolisetty requested review from a team as code owners March 31, 2026 19:28

blunderbuss-gcf Bot assigned duwenxin99 Mar 31, 2026

gemini-code-assist Bot reviewed Mar 31, 2026

View reviewed changes

Comment thread internal/server/server.go

Comment thread internal/sources/cockroachdb/cockroachdb.go

Comment thread internal/sqlcommenter/commenter.go

tanujbolisetty force-pushed the feat/sqlcommenter-integration branch 2 times, most recently from 9836bd1 to b24f64f Compare March 31, 2026 19:35

tanujbolisetty added 5 commits April 10, 2026 10:28

feat: add SQLCommenter observability to database drivers (fixes googl…

1c2d286

…eapis#2899)

fix(mcp): resolve sse memory leak by integrating map cleanup with sse…

ff2d6e2

…Manager

refactor(sources): deduplicate identical RunSQL patterns across pgx a…

10b26f7

…bstraction drivers

fix(sqlcommenter): enforce rfc 3986 percent encoding for spaces

c2ea6d7

chore(license): update copyright headers to 2026 and fix missing license

381e0bc

tanujbolisetty force-pushed the feat/sqlcommenter-integration branch from 3760269 to 09dcc7f Compare April 10, 2026 14:47

tanujbolisetty requested review from a team as code owners April 10, 2026 14:47

docs(sqlcommenter): document session cleanup logic and align module f…

4102ec5

…ramework defaults

tanujbolisetty force-pushed the feat/sqlcommenter-integration branch from 90dbf31 to 4102ec5 Compare April 10, 2026 16:18

averikitsch assigned Yuan325 and unassigned duwenxin99 Apr 13, 2026

Merge branch 'googleapis:main' into feat/sqlcommenter-integration

d520d1a

Yuan325 requested changes Apr 14, 2026

View reviewed changes

tanujbolisetty closed this Apr 26, 2026

tanujbolisetty deleted the feat/sqlcommenter-integration branch April 26, 2026 19:42

		func (s *Source) RunSQL(ctx context.Context, statement string, params parameters.ParamValues, isQuery bool, timeout string) (any, error) {
		// Dgraph does not support SQL block comments, so we explicitly skip injecting SQLCommenter metadata here.

Conversation

tanujbolisetty commented Mar 31, 2026

🚀 Overview

🛠 Features & Modifications

1. Comprehensive Driver Coverage

2. Context Injection Pipeline

3. Core Framework & Build Architecture Upgrades

🧪 Verification

📝 Documentation

Uh oh!

google-cla Bot commented Mar 31, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

drstrangelooker commented Apr 10, 2026

Uh oh!

tanujbolisetty commented Apr 10, 2026

Uh oh!

Yuan325 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Yuan325 left a comment •

edited

Loading